{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f6eb0591-e914-45c9-bf9a-b8f985c8367c",
   "metadata": {},
   "source": [
    "## Usage for MemoRAG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fac7f47b-9ba2-4315-9538-49c6e8c088fd",
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# initialize MemoRAG\n",
    "\n",
    "from memorag import MemoRAG\n",
    "\n",
    "pipe = MemoRAG(\n",
    "    mem_model_name_or_path=\"TommyChien/memorag-mistral-7b-inst\",\n",
    "    ret_model_name_or_path=\"BAAI/bge-m3\",\n",
    "    gen_model_name_or_path=\"mistralai/Mistral-7B-Instruct-v0.2\",\n",
    "    cache_dir=\"path_to_model_cache\",  # to specify local model cache directory (optional)\n",
    "    access_token=\"hugging_face_access_token\"  # to specify local model cache directory (optional)\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20689de7-6714-4f08-b37d-3c14e034d82d",
   "metadata": {},
   "source": [
    "This code block initializes the MemoRAG pipeline using specific model paths for memory (mem_model_name_or_path), retrieval (ret_model_name_or_path), and generation (gen_model_name_or_path). It also sets a cache directory and provides an access token to authenticate the usage of models from Hugging Face or another model-sharing platform."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94516881-5635-4da9-8cbb-5b24a2f23ddc",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Initialize Memory\n",
    "This code block demonstrates how to use the `memorize` function from the MemoRAG pipeline. It begins by loading the contents of a text file (in this case, `harry_potter.txt`), processes the text by memorizing it, and stores the results in a specified directory. In this directory, three key files are created:\n",
    "\n",
    "- **memory.bin**: This file stores the key-value (KV) cache of the memory model, which enables fast retrieval of previously processed information.\n",
    "- **index.bin**: This file contains the dense embeddings for the text corpus, facilitating efficient retrieval of relevant passages.\n",
    "- **chunks.json**: This file holds the passages or chunks derived from the input context, which are used during retrieval.\n",
    "\n",
    "If the `save_dir` parameter is set, the method saves the preprocessed data (i.e., memory, embeddings, and chunks) to disk. This allows for much faster future operations on the same context, as loading the cached data from disk is significantly more efficient than reprocessing and encoding the context from scratch. This caching mechanism is particularly useful when working with large texts or datasets that are frequently accessed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a0f7edb5-5fa6-4a6a-93e9-a396fafd4de9",
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "execution": {
     "iopub.execute_input": "2024-09-05T08:15:49.869863Z",
     "iopub.status.busy": "2024-09-05T08:15:49.869269Z",
     "iopub.status.idle": "2024-09-05T08:17:16.218791Z",
     "shell.execute_reply": "2024-09-05T08:17:16.218175Z",
     "shell.execute_reply.started": "2024-09-05T08:15:49.869839Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory file size: 4.77 GB\n",
      "Encoded context length: 122591 tokens\n",
      "Number of chunks in retrieval corpus: 268\n",
      "Prefilling takes 86.35 second for the full book.\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "start = time.time()\n",
    "test_txt = open(\"harry_potter.txt\").read()\n",
    "pipe.memorize(test_txt, save_dir=\"cache/harry_potter_mistral/\", print_stats=True)\n",
    "print(f\"Prefilling takes {round(time.time()-start,2)} second for the full book.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "62516c2a-7509-440e-88cb-2fff55df4197",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-05T08:17:16.220317Z",
     "iopub.status.busy": "2024-09-05T08:17:16.219813Z",
     "iopub.status.idle": "2024-09-05T08:17:20.057215Z",
     "shell.execute_reply": "2024-09-05T08:17:20.056525Z",
     "shell.execute_reply.started": "2024-09-05T08:17:16.220296Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory file size: 4.77 GB\n",
      "Number of chunks in retrieval corpus: 268\n",
      "Loading from cache takes 3.83 for the full book.\n"
     ]
    }
   ],
   "source": [
    "start = time.time()\n",
    "test_txt = open(\"harry_potter.txt\").read()\n",
    "pipe.load(\"cache/harry_potter_mistral/\", print_stats=True)\n",
    "print(f\"Loading from cache takes {round(time.time()-start,2)} for the full book.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28bfe89a-920e-4cde-8fd0-b1e7727aefda",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-05T06:32:49.870329Z",
     "iopub.status.busy": "2024-09-05T06:32:49.869799Z",
     "iopub.status.idle": "2024-09-05T06:32:49.873295Z",
     "shell.execute_reply": "2024-09-05T06:32:49.872577Z",
     "shell.execute_reply.started": "2024-09-05T06:32:49.870307Z"
    },
    "tags": []
   },
   "source": [
    "### Performing tasks\n",
    "Currently, MemoRAG primarily focuses on two key tasks: question-answering (QA) and summarization.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0d72d41d-e051-44e0-88fb-ed755d98bcb6",
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "execution": {
     "iopub.execute_input": "2024-09-05T08:21:45.739336Z",
     "iopub.status.busy": "2024-09-05T08:21:45.738799Z",
     "iopub.status.idle": "2024-09-05T08:22:43.547922Z",
     "shell.execute_reply": "2024-09-05T08:22:43.547213Z",
     "shell.execute_reply.started": "2024-09-05T08:21:45.739313Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using memory to produce the answer: \n",
      "The article does not provide information on how many times the Chamber of Secrets has been opened in Harry Potter's case. \n",
      "\n",
      "\n",
      "Using MemoRAG to produce the answer: \n",
      "The Chamber of Secrets was opened at least twice in Harry Potter's world, once fifty years prior when a monster attacked students, killing one, and again during Harry's time at Hogwarts.\n"
     ]
    }
   ],
   "source": [
    "query = \"how many times does the chamber be opened in Harry Potter?\"\n",
    "res = pipe(context=test_txt, query=query, task_type=\"qa\", max_new_tokens=256)\n",
    "print(f\"Using memory to produce the answer: \\n{res} \\n\\n\")\n",
    "res = pipe(context=test_txt, query=query, task_type=\"memorag\", max_new_tokens=256)\n",
    "print(f\"Using MemoRAG to produce the answer: \\n{res}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d691689-713a-4794-af51-db58caf1ce95",
   "metadata": {},
   "source": [
    "As demonstrated in the previous examples, relying solely on memory to answer a query can result in an inferior response. This is because the memory serves as a compact and somewhat imprecise representation of the context, leading to a less accurate response. However, when utilizing MemoRAG, the memory model is able to recall key answer clues that guide the retriever to locate more relevant and precise evidence from the original context. This results in a much higher quality response, as the retrieved evidence is more directly aligned with the query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "84577449-1158-4bf5-950f-35df102d400e",
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "execution": {
     "iopub.execute_input": "2024-09-05T08:22:47.202875Z",
     "iopub.status.busy": "2024-09-05T08:22:47.202062Z",
     "iopub.status.idle": "2024-09-05T08:23:53.439885Z",
     "shell.execute_reply": "2024-09-05T08:23:53.439310Z",
     "shell.execute_reply.started": "2024-09-05T08:22:47.202850Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using MemoRAG to summarize the full book:\n",
      " In \"Harry Potter and the Chamber of Secrets\" by J.K. Rowling, Harry Potter is spending his summer break at the Dursleys' house, feeling isolated and unwanted. On his twelfth birthday, he is ignored by his family and longs to return to Hogwarts. One day, while in the garden, Harry encounters a house-elf named Dobby who warns him not to go back to school as there is a plot to cause terrible things at Hogwarts. Harry is skeptical but Dobby shows him a wad of his unopened letters from friends Ron and Hermione, which Dobby had been keeping to prevent Harry from feeling forgotten. Harry becomes angry and insists on reading his letters, causing Dobby to flee.\n",
      "\n",
      "Harry manages to sneak away from the Dursleys and meets Ron and Hermione at the Burrow, the Weasley family home. They plan to visit Hagrid, the Hogwarts gamekeeper, on the weekend. However, Harry is haunted by the memory of Tom Riddle, the past Dark Lord, who had left a diary with Harry and had tried to kill him the previous year. Harry is determined to confront Riddle and find out who or what is behind the attacks at Hogwarts.\n",
      "\n",
      "With the help of Hermione, Harry and Ron come up with a plan to infiltrate the Slytherin common room and interrogate Malfoy about the Chamber of Secrets. They need Polyjuice Potion to transform into Crabbe and Goyle, Malfoy's friends, in order to gain their trust. They manage to obtain the necessary ingredients and brew the potion, but their plan is interrupted by the arrival of Professor Snape.\n",
      "\n",
      "Harry, Ron, and Hermione sneak out of the castle using Harry's Invisibility Cloak and make their way to Hagrid's hut. They find Hagrid guarding the entrance with a crossbow, but he lets them in once he recognizes them. They discuss their plan to find the Chamber of Secrets and confront the monster responsible for the attacks.\n",
      "\n",
      "As they leave Hagrid's hut, they are ambushed by Aragog, a giant Acromantula spider, and his offspring. Harry and Ron are captured and taken to the spider'\n"
     ]
    }
   ],
   "source": [
    "res = pipe(context=test_txt, task_type=\"summarize\", max_new_tokens=512)\n",
    "print(f\"Using MemoRAG to summarize the full book:\\n {res}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f922f083-b956-4d0b-8a43-8558d7a25d6d",
   "metadata": {},
   "source": [
    "### Using APIs as generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0af1eb90-49a7-4e47-b634-bdb6a624f5aa",
   "metadata": {
    "ExecutionIndicator": {
     "show": false
    },
    "execution": {
     "iopub.execute_input": "2024-09-05T08:47:23.769093Z",
     "iopub.status.busy": "2024-09-05T08:47:23.768933Z",
     "iopub.status.idle": "2024-09-05T08:47:46.000519Z",
     "shell.execute_reply": "2024-09-05T08:47:45.999926Z",
     "shell.execute_reply.started": "2024-09-05T08:47:23.769075Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You are using gpt-35-turbo-16k from azure\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Hello! How can I assist you today?']\n",
      "[2024-09-05 16:47:28,238] [INFO] [real_accelerator.py:161:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading checkpoint shards: 100%|██████████| 4/4 [00:09<00:00,  2.47s/it]\n",
      "/opt/conda/lib/python3.10/site-packages/torch/_utils.py:831: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  return self.fget.__get__(instance, owner)()\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory file size: 1.66 GB\n",
      "Number of chunks in retrieval corpus: 268\n"
     ]
    }
   ],
   "source": [
    "from memorag import Agent, MemoRAG\n",
    "api_dict = {\n",
    "    \"endpoint\": \"\",\n",
    "    \"api_version\": \"2024-02-15-preview\",\n",
    "    \"api_key\": \"\"\n",
    "}\n",
    "model = \"gpt-35-turbo-16k\"\n",
    "source = \"azure\"\n",
    "\n",
    "## using deepseek models\n",
    "# model = \"\"\n",
    "# source = \"deepseek\"\n",
    "# api_dict = {\n",
    "#     \"base_url\": \"\",\n",
    "#     \"api_key\": \"\"\n",
    "# }\n",
    "\n",
    "## using openai models#\n",
    "# model = \"\"\n",
    "# source = \"openai\"\n",
    "# api_dict = {\n",
    "#     \"api_key\": \"\"\n",
    "# }\n",
    "\n",
    "\n",
    "agent = Agent(model, source, api_dict)\n",
    "print(agent.generate(\"hi!\")) #  test API\n",
    "\n",
    "pipe = MemoRAG(\n",
    "    mem_model_name_or_path=\"TommyChien/memorag-qwen2-7b-inst\",\n",
    "    ret_model_name_or_path=\"BAAI/bge-m3\",\n",
    "    cache_dir=\"path_to_model_cache\",  # to specify local model cache directory (optional)\n",
    "    customized_gen_model=agent,\n",
    ")\n",
    "pipe.load(\"cache/harry_potter_qwen/\", print_stats=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c0b590b2-c10b-4bd8-97a2-b8617dc0ed6c",
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "execution": {
     "iopub.execute_input": "2024-09-05T08:47:46.003704Z",
     "iopub.status.busy": "2024-09-05T08:47:46.003522Z",
     "iopub.status.idle": "2024-09-05T08:48:05.950280Z",
     "shell.execute_reply": "2024-09-05T08:48:05.949713Z",
     "shell.execute_reply.started": "2024-09-05T08:47:46.003687Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/transformers/generation/configuration_utils.py:427: UserWarning: `do_sample` is set to `False`. However, `top_k` is set to `20` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `top_k`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using MemoRAG with GPT-3.5 to produce the answer: \n",
      "The mutual relationships between the main characters are supportive and loyal. They care for each other and work together to solve problems and overcome challenges.\n"
     ]
    }
   ],
   "source": [
    "query = \"How are the mutual relationships between the main characters? \"\n",
    "test_txt = open(\"harry_potter.txt\").read()\n",
    "\n",
    "res = pipe(context=test_txt, query=query, task_type=\"memorag\", max_new_tokens=256)\n",
    "print(f\"Using MemoRAG with GPT-3.5 to produce the answer: \\n{res}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dae6a550-ac4e-40d5-87be-0c931dd89a0f",
   "metadata": {},
   "source": [
    "## Usage for Memory model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "03d9c8a1-7665-4f55-bed6-ffb26f762e02",
   "metadata": {
    "ExecutionIndicator": {
     "show": false
    },
    "execution": {
     "iopub.execute_input": "2024-09-05T06:55:16.701206Z",
     "iopub.status.busy": "2024-09-05T06:55:16.700649Z",
     "iopub.status.idle": "2024-09-05T06:55:28.069082Z",
     "shell.execute_reply": "2024-09-05T06:55:28.068515Z",
     "shell.execute_reply.started": "2024-09-05T06:55:16.701182Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "Loading checkpoint shards: 100%|██████████| 4/4 [00:10<00:00,  2.66s/it]\n"
     ]
    }
   ],
   "source": [
    "from memorag import Memory\n",
    "\n",
    "memo_model = Memory(\n",
    "    \"TommyChien/memorag-qwen2-7b-inst\",\n",
    "    cache_dir=\"path_to_model_cache\",  # to specify local model cache directory (optional)\n",
    "    beacon_ratio=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "cb319c1a-8fea-487f-98a5-13ef7effa4b3",
   "metadata": {
    "ExecutionIndicator": {
     "show": false
    },
    "execution": {
     "iopub.execute_input": "2024-09-05T06:55:46.665463Z",
     "iopub.status.busy": "2024-09-05T06:55:46.664884Z",
     "iopub.status.idle": "2024-09-05T06:56:06.731255Z",
     "shell.execute_reply": "2024-09-05T06:56:06.730677Z",
     "shell.execute_reply.started": "2024-09-05T06:55:46.665437Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prefilling takes 20.06 second for the full book.\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "start = time.time()\n",
    "context = open(\"harry_potter.txt\").read()\n",
    "memo_model.memorize(context)\n",
    "memo_model.save(\"cache/harry_potter_qwen/memory.bin\")\n",
    "print(f\"Prefilling takes {round(time.time()-start,2)} second for the full book.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "729f7106-009c-4d5e-80ed-ee973db19884",
   "metadata": {
    "ExecutionIndicator": {
     "show": false
    },
    "execution": {
     "iopub.execute_input": "2024-09-05T06:56:06.732919Z",
     "iopub.status.busy": "2024-09-05T06:56:06.732235Z",
     "iopub.status.idle": "2024-09-05T06:56:07.537205Z",
     "shell.execute_reply": "2024-09-05T06:56:07.536655Z",
     "shell.execute_reply.started": "2024-09-05T06:56:06.732899Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading from cache takes 0.8 for the full book.\n"
     ]
    }
   ],
   "source": [
    "memo_model.reset()  # delete memory\n",
    "start = time.time()\n",
    "memo_model.load(\"cache/harry_potter_qwen/memory.bin\")\n",
    "print(f\"Loading from cache takes {round(time.time()-start,2)} for the full book.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "87dad537-cb8b-4e13-8e68-0ebe29bd5596",
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "execution": {
     "iopub.execute_input": "2024-09-05T06:57:15.426928Z",
     "iopub.status.busy": "2024-09-05T06:57:15.426330Z",
     "iopub.status.idle": "2024-09-05T06:57:18.996358Z",
     "shell.execute_reply": "2024-09-05T06:57:18.995804Z",
     "shell.execute_reply.started": "2024-09-05T06:57:15.426902Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using memory to answer the query:\n",
      " Harry, Ron, and Hermione have strong friendships, with Harry and Ron being particularly close. They support each other through difficult times, such as when they are trying to solve the mystery of the Chamber of Secrets. Hermione is also shown to be fiercely loyal to Harry and Ron, often going out of her way to help them.\n"
     ]
    }
   ],
   "source": [
    "query = \"How are the mutual relationships between the main characters? \"\n",
    "\n",
    "res = memo_model.answer(query)\n",
    "print(\"Using memory to answer the query:\\n\", res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "630d527f-4a7d-4e04-8af6-b467b0d1ddd0",
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "execution": {
     "iopub.execute_input": "2024-09-05T07:49:46.166602Z",
     "iopub.status.busy": "2024-09-05T07:49:46.166020Z",
     "iopub.status.idle": "2024-09-05T07:49:59.213827Z",
     "shell.execute_reply": "2024-09-05T07:49:59.213112Z",
     "shell.execute_reply.started": "2024-09-05T07:49:46.166578Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using memory to recall text clues to support the evidence retrieval:\n",
      " 1: Harry Potter and Ron Weasley are best friends. They have been through many adventures together and are always there for each other.\n",
      "2: Hermione Granger is Harry's best friend. She has helped him countless times and they share a strong bond.\n",
      "3: Ron Weasley is Harry's cousin and they are like brothers. They have a close relationship and often tease each other.\n",
      "4: Hogwarts is their home away from home and they miss it dearly when they are apart.\n",
      "5: Harry Potter, Ron Weasley, and Hermione Granger are the main trio in the Harry Potter series. They are loyal to each other and often work together to solve problems.\n",
      "6: Harry Potter and Ron Weasley are like brothers. They have a close relationship and often tease each other.\n",
      "7: Hermione Granger is Harry's best friend. She has helped him countless times and they share a strong bond.\n",
      "8: Ron Weasley is Harry's cousin and they are like brothers. They have a close relationship and often tease each other.\n",
      "9: Harry Potter, Ron Weasley, and Hermione Granger are the main trio in the Harry Potter series. They are loyal to each other and often work together to solve problems.\n",
      "10: Harry Potter and Ron Weasley are like brothers. They have a close relationship and often tease each other.\n"
     ]
    }
   ],
   "source": [
    "res = memo_model.recall(query)\n",
    "res = [line for line in res.split(\"\\n\")[:-1] if line]\n",
    "res = [f\"{i+1}: {line}\" for i,line in enumerate(res)]\n",
    "res = \"\\n\".join(res)\n",
    "print(\"Using memory to recall text clues to support the evidence retrieval:\\n\", res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "a45699cf-f5d3-4f5c-b450-82ee16532ed8",
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "execution": {
     "iopub.execute_input": "2024-09-05T07:06:00.588901Z",
     "iopub.status.busy": "2024-09-05T07:06:00.588323Z",
     "iopub.status.idle": "2024-09-05T07:06:04.146577Z",
     "shell.execute_reply": "2024-09-05T07:06:04.145964Z",
     "shell.execute_reply.started": "2024-09-05T07:06:00.588878Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using memory to rewrite the input query into more specifc surrogate queries:\n",
      " 1: What are the interactions and relationships between Harry Potter, Ron Weasley, Hermione Granger, and Draco Malfoy?\n",
      "2: How do Harry Potter and Ron Weasley support each other during challenging situations?\n",
      "3: What role does Hermione Granger play in solving mysteries and puzzles?\n"
     ]
    }
   ],
   "source": [
    "res = memo_model.rewrite(query)\n",
    "res = [f\"{i+1}: {line}\" for i,line in enumerate(res.split(\"\\n\")[:-1]) if line]\n",
    "res = \"\\n\".join(res)\n",
    "print(\"Using memory to rewrite the input query into more specifc surrogate queries:\\n\", res)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb7890fa-51ec-4b72-9a53-d33121976a92",
   "metadata": {},
   "source": [
    "## Usage for Memory-augmented retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5b3131d-3c99-484b-a45c-b0ff5c55e143",
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from memorag import MemoRAG\n",
    "\n",
    "pipe = MemoRAG(\n",
    "    mem_model_name_or_path=\"TommyChien/memorag-qwen2-7b-inst\",\n",
    "    ret_model_name_or_path=\"BAAI/bge-m3\",\n",
    "    cache_dir=\"path_to_model_cache\",  # to specify local model cache directory (optional)\n",
    "    access_token=\"hugging_face_access_token\"  # to specify local model cache directory (optional)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c549d2a3-e28b-45c9-a261-cbbbaa4d4043",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-05T08:25:15.989351Z",
     "iopub.status.busy": "2024-09-05T08:25:15.988610Z",
     "iopub.status.idle": "2024-09-05T08:25:39.837902Z",
     "shell.execute_reply": "2024-09-05T08:25:39.837283Z",
     "shell.execute_reply.started": "2024-09-05T08:25:15.989329Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory file size: 1.66 GB\n",
      "Encoded context length: 122591 tokens\n",
      "Number of chunks in retrieval corpus: 268\n",
      "Prefilling takes 23.85 second for the full book.\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "start = time.time()\n",
    "test_txt = open(\"harry_potter.txt\").read()\n",
    "pipe.memorize(test_txt, save_dir=\"cache/harry_potter_qwen/\", print_stats=True)\n",
    "print(f\"Prefilling takes {round(time.time()-start,2)} second for the full book.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e9c29f68-f5d3-41e5-82be-2a092030fa62",
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "execution": {
     "iopub.execute_input": "2024-09-05T08:26:29.273539Z",
     "iopub.status.busy": "2024-09-05T08:26:29.272717Z",
     "iopub.status.idle": "2024-09-05T08:26:30.578458Z",
     "shell.execute_reply": "2024-09-05T08:26:30.577911Z",
     "shell.execute_reply.started": "2024-09-05T08:26:29.273505Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory file size: 1.66 GB\n",
      "Number of chunks in retrieval corpus: 268\n",
      "Loading from cache takes 1.3 for the full book.\n"
     ]
    }
   ],
   "source": [
    "start = time.time()\n",
    "test_txt = open(\"harry_potter.txt\").read()\n",
    "pipe.load(\"cache/harry_potter_qwen/\", print_stats=True)\n",
    "print(f\"Loading from cache takes {round(time.time()-start,2)} for the full book.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "40bd5a0d-2c97-40fb-b1f4-b7eeedd7253c",
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "execution": {
     "iopub.execute_input": "2024-09-05T08:28:39.991651Z",
     "iopub.status.busy": "2024-09-05T08:28:39.990931Z",
     "iopub.status.idle": "2024-09-05T08:28:53.582693Z",
     "shell.execute_reply": "2024-09-05T08:28:53.581954Z",
     "shell.execute_reply.started": "2024-09-05T08:28:39.991627Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/transformers/generation/configuration_utils.py:427: UserWarning: `do_sample` is set to `False`. However, `top_k` is set to `20` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `top_k`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Harry Potter and Ron Weasley are best friends. They have been through many adventures together and are always there for each other.', \"Hermione Granger is Harry's best friend. She has helped him countless times and they share a strong bond.\", \"Ron Weasley is Harry's cousin and they are like brothers. They have a close relationship and often tease each other.\", 'Hogwarts is their home away from home and they miss it dearly when they are apart.', 'Harry Potter, Ron Weasley, and Hermione Granger are the main trio in the Harry Potter series. They are loyal to each other and often work together to solve problems.', 'Harry Potter and Ron Weasley are like brothers. They have a close relationship and often tease each other.', \"Hermione Granger is Harry's best friend. She has helped him countless times and they share a strong bond.\", \"Ron Weasley is Harry's cousin and they are like brothers. They have a close relationship and often tease each other.\", 'Harry Potter, Ron Weasley, and Hermione Granger are the main trio in the Harry Potter series. They are loyal to each other and often work together to solve problems.', 'Harry Potter and Ron Weasley are like brothers. They have a close relationship and often tease each other.']\n"
     ]
    }
   ],
   "source": [
    "query = \"How are the mutual relationships between the main characters? \"\n",
    "\n",
    "clues = pipe.mem_model.recall(query).split(\"\\n\")\n",
    "clues = [q for q in clues if len(q.split()) > 3]\n",
    "print(clues)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f31a356a-25e7-401c-86fb-dbd6162d6b7b",
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "execution": {
     "iopub.execute_input": "2024-09-05T08:30:41.735218Z",
     "iopub.status.busy": "2024-09-05T08:30:41.734686Z",
     "iopub.status.idle": "2024-09-05T08:30:41.755973Z",
     "shell.execute_reply": "2024-09-05T08:30:41.755352Z",
     "shell.execute_reply.started": "2024-09-05T08:30:41.735197Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "He missed Hogwarts so much it was like having a constant\n",
      "stomachache. He missed the castle, with its secret passageways and\n",
      "ghosts, his classes (though perhaps not Snape, the Potions master), the\n",
      "mail arriving by owl, eating banquets in the Great Hall, sleeping in his\n",
      "four-poster bed in the tower dormitory, visiting the gamekeeper,\n",
      "Hagrid, in his cabin next to the Forbidden Forest in the grounds, and,\n",
      "especially, Quidditch, the most popular sport in the wizarding world\n",
      "(six tall goal posts, four flying balls, and fourteen players on\n",
      "broomsticks).\n",
      "\n",
      "All Harry's spellbooks, his wand, robes, cauldron, and top-of-the-line\n",
      "Nimbus Two Thousand broomstick had been locked in a cupboard\n",
      "under the stairs by Uncle Vernon the instant Harry had come home.\n",
      "What did the Dursleys care if Harry lost his place on the House\n",
      "Quidditch team because he hadn't practiced all summer? What was it\n",
      "to the Dursleys if Harry went back to school without any of his\n",
      "homework done? The Dursleys were what wizards called Muggles\n",
      "(not a drop of magical blood in their veins),\n",
      "\n",
      "and as far as they were concerned, having a wizard in the family was\n",
      "a matter of deepest shame. Uncle Vernon had even padlocked\n",
      "Harry's owl, Hedwig, inside her cage, to stop her from carrying\n",
      "messages to anyone in the wizarding world.\n",
      "\n",
      "Harry looked nothing like the rest of the family. Uncle Vernon was\n",
      "large and neckless, with an enormous black mustache; Aunt Petunia\n",
      "was horse-faced and bony; Dudley was blond, pink, and porky. Harry,\n",
      "on the other hand, was small and skinny, with brilliant green eyes and\n",
      "jet-black hair that was always untidy. He wore round glasses, and on\n",
      "his forehead was a thin, lightning-shaped scar.\n",
      "\n",
      "It was this scar that made Harry so particularly unusual, even for a\n",
      "wizard. This scar was the only hint of Harry's very mysterious past, of\n",
      "the reason he had been left on the Dursleys' doorstep eleven years\n",
      "before.\n",
      "======\n",
      "Shaking, Harry let Dobby out of the closet.\n",
      "\n",
      "\"See what it's like here?\" he said. \"See why I've got to go back to\n",
      "Hogwarts? It's the only place I've got -well, I think I've got friends. \"\n",
      "\n",
      "\"Friends who don't even write to Harry Potter?\" said Dobby slyly.\n",
      "\n",
      "\"I expect they've just been - wait a minute,\" said Harry, frowning.\n",
      "\"How do you know my friends haven't been writing to me?\"\n",
      "\n",
      "Dobby shuffled his feet.\n",
      "\n",
      "\"Harry Potter mustn't be angry with Dobby. Dobby did it for the best -\n",
      "\"\n",
      "\n",
      "\"Have you been stopping my letters?\"\n",
      "\n",
      "\"Dobby has them here, sir,\" said the elf. Stepping nimbly out of Harry's\n",
      "reach, he pulled a thick wad of envelopes from the inside of the\n",
      "pillowcase he was wearing. Harry could make out Hermione's neat\n",
      "writing, Ron's untidy scrawl, and even a scribble that looked as though\n",
      "it was from the Hogwarts gamekeeper, Hagrid.\n",
      "\n",
      "Dobby blinked anxiously up at Harry.\n",
      "\n",
      "\"Harry Potter mustn't be angry... Dobby hoped ... if Harry Potter\n",
      "thought his friends had forgotten him ... Harry Potter might not want to\n",
      "go back to school, sir . .....\n",
      "\n",
      "Harry wasn't listening. He made a grab for the letters, but Dobby\n",
      "jumped out of reach.\n",
      "\n",
      "\"Harry Potter will have them, sir, if he gives Dobby his word\n",
      "\n",
      "that he will not return to Hogwarts. Ah, sir, this is a danger you must\n",
      "not face! Say you won't go back, sir!\"\n",
      "\n",
      "\"No,\" said Harry angrily. \"Give me my friends' letters!\"\n",
      "\n",
      "\"Then Harry Potter leaves Dobby no choice,\" said the elf sadly.\n",
      "\n",
      "Before Harry could move, Dobby had darted to the bedroom door,\n",
      "pulled it open, and sprinted down the stairs.\n",
      "\n",
      "Mouth dry, stomach lurching, Harry sprang after him, trying not to\n",
      "make a sound. He jumped the last six steps, landing catlike on the\n",
      "hall carpet, looking around for Dobby. From the dining room he\n",
      "heard Uncle Vernon saying, \". . . tell Petunia that very funny story\n",
      "about those American plumbers, Mr. Mason. She's been dying to\n",
      "hear. . . \"\n",
      "\n",
      "Harry ran up the hall into the kitchen and felt his stomach disappear.\n",
      "======\n",
      "He dreamed that he was on show in a zoo, with a card reading\n",
      "UNDERAGE WIZARD attached to his cage. People goggled through the bars\n",
      "at him as he lay, starving and weak, on a bed of straw. He saw\n",
      "Dobby's face in the crowd and shouted out, asking for help, but Dobby\n",
      "called, \"Harry Potter is safe there, sir!\" and vanished. Then the\n",
      "Dursleys appeared and Dudley rattled the bars of the cage, laughing at\n",
      "him.\n",
      "\n",
      "\"Stop it,\" Harry muttered as the rattling pounded in his sore head.\n",
      "\"Leave me alone ... cut it out ... I'm trying to sleep . . . .\"\n",
      "\n",
      "He opened his eyes. Moonlight was shining through the bars on the\n",
      "window. And someone was goggling through the bars at him: a freckle-\n",
      "faced, red-haired, long-nosed someone.\n",
      "\n",
      "Ron Weasley was outside Harry's window.\n",
      "\n",
      "H-H A P T E RR T 11-H RR E E\n",
      "\n",
      "THE BURROW\n",
      "\n",
      "Ron.l\" breathed Harry, creeping to the window and pushing it up so\n",
      "they could talk through the bars. \"Ron, how did you - What the -?\"\n",
      "\n",
      "Harry's mouth fell open as the full impact of what he was seeing hit\n",
      "him. Ron was leaning out of the back window of an old turquoise car,\n",
      "which was parked in midair Grinning at Harry from the front seats\n",
      "were Fred and George, Ron's elder twin brothers.\n",
      "\n",
      "\"All right, Harry?\" asked George.\n",
      "\n",
      "\"What's been going on?\" said Ron. \"Why haven't you been answering\n",
      "my letters? I've asked you to stay about twelve times, and then Dad\n",
      "came home and said you'd got an official warning for using magic in\n",
      "front of Muggles -\"\n",
      "\n",
      "\"It wasn't me - and how did he know?\"\n",
      "\n",
      "\"He works for the Ministry,\" said Ron. \"You know we're not supposed\n",
      "to do spells outside school -\"\n"
     ]
    }
   ],
   "source": [
    "retrieved_passages = pipe._retrieve(clues)\n",
    "print(\"\\n======\\n\".join(retrieved_passages[:3]))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
